
from monai.networks.nets.basic_unet import BasicUNet
import torch.nn as nn
from medical.loss.dice_ce_loss import EntropyAndDiceLoss

class UNet(nn.Module):
    def __init__(self,
                 in_channels,
                 out_channels,
                 features=(32, 32, 64, 128, 256, 32),
                 dropout=0.0
                 ):
        super().__init__()
        self.model = BasicUNet(spatial_dims=3,
                               in_channels=in_channels,
                               out_channels=out_channels,
                               features=features,
                               dropout=dropout)

        self.dice_ce_loss = EntropyAndDiceLoss()

    def forward(self, x):
        if type(x) is dict:
            return self.forward_train(x)

        return self.model(x)

    def forward_train(self, data):
        x = data["image"]
        label = data.get("label", None)

        logits = self.forward(x)

        return_data = {"logits": logits}
        if label is not None:
            loss = self.compute_loss(logits, label)
            return_data["loss"] = loss

        return return_data

    def compute_loss(self, pred, label):
        return self.dice_ce_loss(pred, label)